import argparse
from models import *  # set ONNX_EXPORT in models.py


def zjm_fuse():
    # Initialize model
    model = Darknet(opt.cfg_in, opt.img_size)

    # Load weights
    if opt.weight_in.endswith('.pt'):  # pytorch format
        model.load_state_dict(torch.load(opt.weight_in)['model'])
    elif len(opt.weight_in) > 0:  # darknet format
        cutoff = load_darknet_weights(model, opt.weight_in)
        
    model.fuse()
    #############
    
    rebuildnet = Darknet(opt.cfg_out, opt.img_size)
    netwgt_list = [key for key in model.state_dict().keys()]
    rebuildnetwgt_list = [key for key in rebuildnet.state_dict().keys()]
    # print(rebuildnetwgt_list)
    renet2net_dict = dict()
    for i in range(len(netwgt_list)):
        renet2net_dict[rebuildnetwgt_list[i]] = netwgt_list[i]
    
    rebuildnet_state_dict = rebuildnet.state_dict()
    for key in rebuildnet.state_dict().keys():
        rebuildnet_state_dict[key] = model.state_dict()[renet2net_dict[key]]
    
    rebuildnet.load_state_dict(rebuildnet_state_dict)
    
    #############
    torch.save({'model' : rebuildnet.state_dict()}, opt.weight_out)
    # save_weights(rebuildnet, opt.weight_out)
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg-in', type=str, default='./cfg/yolov3_ReLU-voc.cfg', help='input *.cfg path')
    parser.add_argument('--cfg-out', type=str, default='./cfg/yolov3_ReLU_fuse-voc.cfg', help='output *.cfg path')
    parser.add_argument('--weight-in', type=str, default='/home/zhangjm/backup/yolov3_ReLU_voc/yolov3_ReLU-voc_final.weights', help='path to input weights file')
    parser.add_argument('--weight-out', type=str, default='/home/zhangjm/backup/yolov3_ReLU_fuse_voc/20200321.pt', help='output weights')
    parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')

    opt = parser.parse_args()
    print(opt)

    with torch.no_grad():
        zjm_fuse()
